import os
import gradio as gr     #使用gradio框架，用于构建前端
import shutil     #进行对文件的删除，移动等操作
import jieba.analyse as aly  #这是一个中文分词组件（结巴分词）
from collections import Counter    #用于统计，如aaabbcc  则a:3 b:2 c:2
from database import MyDataBase    #导入自己的向量数据库
from database import load_database
import pandas as pd
from local_llm_model import get_ans   #导入大模型，根据提示得到答案
from knowledge_extract import generate_qa   #知识库的qa问答对的提取


def get_type_name(files):
    content = []
    for file in files:
        try:
            with open(file.name, encoding="utf-8") as f:
                data = f.readlines(1)
                content.extend(aly.tfidf(data[0]))
        except:
            continue
    count = Counter(content)
    kw = count.most_common(2)

    return "".join([i[0] for i in kw])


def upload(files):
    global database_list, database_namelist, input_qa_gen,input_qa_ans
    #传入一个知识库，里面包含多个txt文件
    check_txt = False
    '''
    for else用法，如果for正常结束，则执行else块中的内容。
    非正常结束，被break终端掉，则不执行else中的内容
    '''
    for file in files:
        if check_txt:
            break
        if file.name.endswith(".txt"):
            check_txt = True
    else:
        #如果文件夹下没有txt文档，则报错。只要有一个txt文档，就会中断上面的for，从而正常读取
        if check_txt == False:
            raise Exception("请上传包含txt文档的文件夹")

    type_name = get_type_name(files)  #这个就是为了得到当前数据库的名称，其中统计频率最高的几个词作为数据库的名称
    save_path = os.path.join("data/database_dir", type_name) #知识库存储路径

    if os.path.exists(save_path) == False:
        os.mkdir(save_path)#没有这个路径，就创建这样一个知识库
        os.mkdir(os.path.join(save_path, "qa")) #创建一个qa文件夹，要不然会报错
        os.mkdir(os.path.join(save_path, "txt")) #再在知识库下创建一个txt路径
    for file in files:
        if file.name.endswith(".txt"):
            shutil.copy(file.name, os.path.join(save_path, "txt"))  #将所有的txt文件复制到这个知识库路径的txt文件夹下

    database = MyDataBase(save_path, type_name)  #创建一个MyDataBase对象
    database_list.append(database)  #mydatabase对象加入到database_list中
    database_namelist.append(type_name)   #再加入知识库的名称

    knowledge_names.choices.append((type_name, type_name))  #上传模块添加该知识库
    input_qa_gen.choices.append((type_name, type_name))     #知识抽取模块添加该知识库
    input_qa_ans.choices.append((type_name, type_name))     #问答模块添加该知识库
    context = pd.DataFrame(database.document.contents, columns=["context"])  #更新知识库内容
    return type_name, context   #返回知识库名称和文件内容


def database_change(name):     #传入一个知识库名称
    global database_list, database_namelist
    #创建一个dataframe对象，可是这里问什么是question?不应该是context？
    context = pd.DataFrame(database_list[database_namelist.index(name)].document.contents, columns=["context"])
    return context


def generate_rag_result(knowledge_name, question):  #传入知识库和问题
    # Call your data generation function here
    global database_list, database_namelist
    database = database_list[database_namelist.index(knowledge_name)]
    search_result = database.search(question, 3)   #同通过该知识库得到最相关的三个qa问答对
    abstract = "\n".join(search_result["answer"])  #得到对应的答案，基于答案，回答用户的问题,每个答案加入一个换行符号
    print(f'abstract-------------->{abstract}')
    #'''{abstract}'''表明其中的多行文本
    prompt = f"请根据已知内容简洁明了的回复学生的问题，已知内容如下：```{abstract}```,用户的问题是：{question}，如何已知内容无法回答用户的问题，请直接回复：不知道，无需输出其他内容"
    result = get_ans(prompt)   #这里就是将答案和提示词去传递给大模型，利用大模型的能力去完成输出
    return result, search_result    #返回的是，答案  和 rag搜寻的3个qa问答对


def process_qa_data(input_qa):     #处理qa数据，接受一个知识库的名字
    global database_list, database_namelist
    select_database = database_list[database_namelist.index(input_qa)]
    contexts = select_database.document.contents  #得到该数据库的所有文档切片
    qa_df = generate_qa(input_qa, contexts)    #进行知识提取，传入知识库的名称和所有的文档切片，返回一个qa问答对的DataFrame
    select_database.create_emb_database(qa_df)  #为该数据库选择本地的embedding模型，并且传入qa问答对（知识提取的结果）
    return qa_df  #返回的是qa问答对


if __name__ == "__main__":

    #默认选择打开data/database_dir
    #database_list是MyDataBase对象 ，database_namelist是数据库的名称
    #比如小米汽车就相当于是一个数据库
    database_list, database_namelist = load_database()
    print("data................")

    with gr.Blocks() as demo:
        with gr.Tab("知识库管理"):
            knowledge_names = gr.Dropdown(choices=database_namelist, label="知识库选择", value=database_namelist[0])#默认选择第一个知识库
            context = gr.DataFrame(pd.DataFrame(database_list[0].document.contents, columns=["context"]), height=800)  #展示文档的内容，传入一个dataFrame
            input3 = gr.UploadButton(label="上传知识库", file_count="directory")  #上传的类型选择目录，有single，multiple ,dreectory
            input3.upload(upload, input3, [knowledge_names, context])  #这里相当于接受一个函数，输入是input3，输出是
            knowledge_names.change(database_change, knowledge_names, context)  #这里增加监听，只要知识库改变，它下方的context就改变
        with gr.Tab("生成问答数据"):
            input_qa_gen = gr.Dropdown(choices=database_namelist, label="知识库选择", value=database_namelist[0])
            output_data = gr.DataFrame(database_list[0].document.qa_df, height=400)  #默认选择第0个的qa问答对
            generate_data_button = gr.Button("生成问答QA数据")
            generate_data_button.click(process_qa_data, [input_qa_gen], [output_data])    #接受一个知识库名称，返回一个dataframe，即知识提取结果

        with gr.Tab("RAG知识库问答"):
            input_qa_ans = gr.Dropdown(choices=database_namelist, label="知识库问答", value=database_namelist[0])
            text_input = gr.Textbox(label="请输出问题")   #输入问题文本框
            text_rag = gr.DataFrame(pd.DataFrame(), height=400,label="RAG 结果")
            text_output = gr.Textbox(label="大模型生成答案")
            qa_button = gr.Button("生成答案")
            qa_button.click(generate_rag_result, [input_qa_ans, text_input], [text_output,text_rag]) #传入知识库和问题

    demo.launch(server_name="127.0.0.1", server_port=9999, show_api=False, auth=("username", "password"))
